import requests
import json
import os
import time
from PIL import Image
import io

base_url = "http://localhost:8080/api"
avatar_url = f"{base_url}/avatar"
users_url = f"{base_url}/users"


# 首先登入用户
def login_user(username, password):
    """用户登录并返回token"""
    login_data = {"username": username, "password": password}
    response = requests.post(
        f"{users_url}/login",
        json=login_data,
        headers={"Content-Type": "application/json"},
    )
    if response.status_code == 200:
        return response.json().get("token")
    else:
        print(f"登录失败: {response.json()}")
        return None


def register_user(username, password):
    """用户注册"""
    register_data = {"username": username, "password": password}
    response = requests.post(
        f"{users_url}/register",
        json=register_data,
        headers={"Content-Type": "application/json"},
    )
    return response.status_code == 200


def create_test_image(filename, format="JPEG", size=(100, 100), color="red"):
    """创建测试图片文件"""
    try:
        # 创建一个纯色图片
        image = Image.new("RGB", size, color)
        
        # 保存为指定格式
        image.save(filename, format)
        print(f"✅ 成功创建测试图片: {filename}")
        return True
    except Exception as e:
        print(f"❌ 创建测试图片失败: {e}")
        return False


def create_large_test_image(filename, size_mb=6):
    """创建超过限制大小的测试图片"""
    try:
        # 计算需要的像素数量来达到指定大小
        target_size = size_mb * 1024 * 1024  # 转换为字节
        # 估算需要的图片尺寸（RGB每像素3字节）
        pixels_needed = target_size // 3
        width = int(pixels_needed ** 0.5)
        height = width
        
        # 创建大图片
        image = Image.new("RGB", (width, height), "blue")
        image.save(filename, "JPEG", quality=100)  # 最高质量确保大小
        
        # 检查文件大小
        file_size = os.path.getsize(filename)
        print(f"✅ 成功创建大图片: {filename}, 大小: {file_size / 1024 / 1024:.2f}MB")
        return True
    except Exception as e:
        print(f"❌ 创建大图片失败: {e}")
        return False


def test_avatar_info():
    """测试获取头像管理接口信息"""
    print("=== 测试获取头像管理接口信息 ===")

    try:
        response = requests.get(f"{avatar_url}/")

        print(f"状态码: {response.status_code}")
        print(f"响应内容: {json.dumps(response.json(), ensure_ascii=False, indent=2)}")

        if response.status_code == 200:
            print("✅ 获取头像接口信息成功")
            return True
        else:
            print("❌ 获取头像接口信息失败")
            return False

    except Exception as e:
        print(f"❌ 获取头像接口信息请求异常: {e}")
        return False


def test_upload_avatar_success(token, image_filename):
    """测试成功上传头像"""
    print(f"\n=== 测试成功上传头像 ({image_filename}) ===")

    if not token:
        print("❌ 没有有效的token，跳过测试")
        return False

    if not os.path.exists(image_filename):
        print(f"❌ 测试图片文件不存在: {image_filename}")
        return False

    headers = {
        "Authorization": f"Bearer {token}",
    }

    try:
        with open(image_filename, 'rb') as f:
            files = {'file': (image_filename, f, 'image/jpeg')}
            response = requests.post(f"{avatar_url}/upload", files=files, headers=headers)

        print(f"状态码: {response.status_code}")
        response_data = response.json()
        print(f"响应内容: {json.dumps(response_data, ensure_ascii=False, indent=2)}")

        if response.status_code == 200:
            avatar_url_returned = response_data.get("avatar_url", "")
            print(f"✅ 头像上传成功，URL: {avatar_url_returned}")
            return True
        else:
            print("❌ 头像上传失败")
            return False

    except Exception as e:
        print(f"❌ 头像上传请求异常: {e}")
        return False


def test_upload_avatar_unauthorized():
    """测试未授权上传头像"""
    print("\n=== 测试未授权上传头像 ===")

    # 创建测试图片
    test_image = "test_unauthorized.jpg"
    create_test_image(test_image)

    try:
        # 测试没有token的情况
        with open(test_image, 'rb') as f:
            files = {'file': (test_image, f, 'image/jpeg')}
            response = requests.post(f"{avatar_url}/upload", files=files)
        print(f"无token访问 - 状态码: {response.status_code}")

        # 测试无效token的情况
        headers = {"Authorization": "Bearer invalid_token"}
        with open(test_image, 'rb') as f:
            files = {'file': (test_image, f, 'image/jpeg')}
            response = requests.post(f"{avatar_url}/upload", files=files, headers=headers)
        print(f"无效token访问 - 状态码: {response.status_code}")

        if response.status_code == 401:
            print("✅ 未授权访问验证成功")
            return True
        else:
            print("❌ 未授权访问验证失败")
            return False

    except Exception as e:
        print(f"❌ 未授权访问测试异常: {e}")
        return False
    finally:
        # 清理测试文件
        if os.path.exists(test_image):
            os.remove(test_image)


def test_upload_invalid_file_format(token):
    """测试上传无效文件格式"""
    print("\n=== 测试上传无效文件格式 ===")

    if not token:
        print("❌ 没有有效的token，跳过测试")
        return False

    headers = {
        "Authorization": f"Bearer {token}",
    }

    # 创建文本文件
    txt_file = "test_invalid.txt"
    try:
        with open(txt_file, 'w') as f:
            f.write("这是一个文本文件，不是图片")

        with open(txt_file, 'rb') as f:
            files = {'file': (txt_file, f, 'text/plain')}
            response = requests.post(f"{avatar_url}/upload", files=files, headers=headers)

        print(f"状态码: {response.status_code}")
        print(f"响应内容: {response.json()}")

        if response.status_code == 400:
            print("✅ 无效文件格式验证成功")
            return True
        else:
            print("❌ 无效文件格式验证失败")
            return False

    except Exception as e:
        print(f"❌ 无效文件格式测试异常: {e}")
        return False
    finally:
        # 清理测试文件
        if os.path.exists(txt_file):
            os.remove(txt_file)


def test_upload_large_file(token):
    """测试上传超大文件"""
    print("\n=== 测试上传超大文件 (>5MB) ===")

    if not token:
        print("❌ 没有有效的token，跳过测试")
        return False

    headers = {
        "Authorization": f"Bearer {token}",
    }

    large_image = "test_large.jpg"
    try:
        # 创建超大图片（6MB）
        if not create_large_test_image(large_image, size_mb=6):
            return False

        with open(large_image, 'rb') as f:
            files = {'file': (large_image, f, 'image/jpeg')}
            response = requests.post(f"{avatar_url}/upload", files=files, headers=headers)

        print(f"状态码: {response.status_code}")
        print(f"响应内容: {response.json()}")

        if response.status_code == 400:
            print("✅ 大文件限制验证成功")
            return True
        else:
            print("❌ 大文件限制验证失败")
            return False

    except Exception as e:
        print(f"❌ 大文件测试异常: {e}")
        return False
    finally:
        # 清理测试文件
        if os.path.exists(large_image):
            os.remove(large_image)


def test_upload_different_formats(token):
    """测试上传不同格式的图片"""
    print("\n=== 测试上传不同格式的图片 ===")

    if not token:
        print("❌ 没有有效的token，跳过测试")
        return False

    headers = {
        "Authorization": f"Bearer {token}",
    }

    formats = [
        ("test_jpeg.jpg", "JPEG", "image/jpeg"),
        ("test_png.png", "PNG", "image/png"),
    ]

    success_count = 0

    for filename, format_name, content_type in formats:
        try:
            print(f"\n--- 测试 {format_name} 格式 ---")
            
            # 创建测试图片
            create_test_image(filename, format=format_name, color="green")

            with open(filename, 'rb') as f:
                files = {'file': (filename, f, content_type)}
                response = requests.post(f"{avatar_url}/upload", files=files, headers=headers)

            print(f"状态码: {response.status_code}")
            print(f"响应内容: {response.json()}")

            if response.status_code == 200:
                print(f"✅ {format_name} 格式上传成功")
                success_count += 1
            else:
                print(f"❌ {format_name} 格式上传失败")

        except Exception as e:
            print(f"❌ {format_name} 格式测试异常: {e}")
        finally:
            # 清理测试文件
            if os.path.exists(filename):
                os.remove(filename)

    print(f"\n✅ 格式测试完成，成功: {success_count}/{len(formats)}")
    return success_count == len(formats)


def test_upload_no_file(token):
    """测试不提供文件的上传请求"""
    print("\n=== 测试不提供文件的上传请求 ===")

    if not token:
        print("❌ 没有有效的token，跳过测试")
        return False

    headers = {
        "Authorization": f"Bearer {token}",
    }

    try:
        # 发送空的文件上传请求
        response = requests.post(f"{avatar_url}/upload", headers=headers)

        print(f"状态码: {response.status_code}")
        print(f"响应内容: {response.json()}")

        if response.status_code == 422:  # FastAPI参数验证错误
            print("✅ 缺少文件验证成功")
            return True
        else:
            print("❌ 缺少文件验证失败")
            return False

    except Exception as e:
        print(f"❌ 缺少文件测试异常: {e}")
        return False


def test_upload_empty_file(token):
    """测试上传空文件"""
    print("\n=== 测试上传空文件 ===")

    if not token:
        print("❌ 没有有效的token，跳过测试")
        return False

    headers = {
        "Authorization": f"Bearer {token}",
    }

    empty_file = "test_empty.jpg"
    try:
        # 创建空文件
        with open(empty_file, 'w') as f:
            pass  # 创建空文件

        with open(empty_file, 'rb') as f:
            files = {'file': (empty_file, f, 'image/jpeg')}
            response = requests.post(f"{avatar_url}/upload", files=files, headers=headers)

        print(f"状态码: {response.status_code}")
        print(f"响应内容: {response.json()}")

        if response.status_code == 400 or response.status_code == 500:
            print("✅ 空文件验证成功")
            return True
        else:
            print("❌ 空文件验证失败")
            return False

    except Exception as e:
        print(f"❌ 空文件测试异常: {e}")
        return False
    finally:
        # 清理测试文件
        if os.path.exists(empty_file):
            os.remove(empty_file)


def test_multiple_uploads(token):
    """测试多次上传头像（覆盖之前的头像）"""
    print("\n=== 测试多次上传头像 ===")

    if not token:
        print("❌ 没有有效的token，跳过测试")
        return False

    headers = {
        "Authorization": f"Bearer {token}",
    }

    success_count = 0
    colors = ["red", "blue", "yellow"]

    for i, color in enumerate(colors, 1):
        filename = f"test_multiple_{i}.jpg"
        try:
            print(f"\n--- 第{i}次上传 ({color}色头像) ---")
            
            # 创建不同颜色的测试图片
            create_test_image(filename, color=color)

            with open(filename, 'rb') as f:
                files = {'file': (filename, f, 'image/jpeg')}
                response = requests.post(f"{avatar_url}/upload", files=files, headers=headers)

            print(f"状态码: {response.status_code}")
            response_data = response.json()
            print(f"响应内容: {json.dumps(response_data, ensure_ascii=False, indent=2)}")

            if response.status_code == 200:
                print(f"✅ 第{i}次上传成功")
                success_count += 1
            else:
                print(f"❌ 第{i}次上传失败")

        except Exception as e:
            print(f"❌ 第{i}次上传异常: {e}")
        finally:
            # 清理测试文件
            if os.path.exists(filename):
                os.remove(filename)

    print(f"\n✅ 多次上传测试完成，成功: {success_count}/{len(colors)}")
    return success_count == len(colors)


def test_concurrent_uploads(token):
    """测试并发上传（简单版本）"""
    print("\n=== 测试并发上传 ===")

    if not token:
        print("❌ 没有有效的token，跳过测试")
        return False

    headers = {
        "Authorization": f"Bearer {token}",
    }

    # 创建测试图片
    test_image = "test_concurrent.jpg"
    create_test_image(test_image, color="purple")

    success_count = 0
    total_requests = 3

    try:
        import threading
        import time

        def upload_avatar(index):
            nonlocal success_count
            try:
                with open(test_image, 'rb') as f:
                    files = {'file': (f"avatar_{index}.jpg", f, 'image/jpeg')}
                    response = requests.post(f"{avatar_url}/upload", files=files, headers=headers)
                
                if response.status_code == 200:
                    success_count += 1
                    print(f"✅ 并发请求 {index} 成功")
                else:
                    print(f"❌ 并发请求 {index} 失败: {response.status_code}")
            except Exception as e:
                print(f"❌ 并发请求 {index} 异常: {e}")

        # 创建并启动线程
        threads = []
        for i in range(total_requests):
            thread = threading.Thread(target=upload_avatar, args=(i+1,))
            threads.append(thread)
            thread.start()

        # 等待所有线程完成
        for thread in threads:
            thread.join()

        print(f"\n✅ 并发上传测试完成，成功: {success_count}/{total_requests}")
        return success_count > 0  # 至少有一个成功

    except ImportError:
        print("❌ 无法导入threading模块，跳过并发测试")
        return True
    except Exception as e:
        print(f"❌ 并发上传测试异常: {e}")
        return False
    finally:
        # 清理测试文件
        if os.path.exists(test_image):
            os.remove(test_image)


def run_all_avatar_tests():
    """运行所有头像测试"""
    print("开始头像功能测试...")
    print("=" * 60)

    # 测试用户名（使用时间戳确保唯一性）
    username = f"avatar_test_user_{int(time.time())}"
    password = "123456"

    # 1. 注册测试用户
    print(f"注册测试用户: {username}")
    if not register_user(username, password):
        print("❌ 用户注册失败，测试终止")
        return

    # 2. 登录获取token
    print(f"登录测试用户: {username}")
    token = login_user(username, password)
    if not token:
        print("❌ 用户登录失败，测试终止")
        return

    print(f"✅ 获得token: {token[:20]}...")

    # 3. 测试获取头像接口信息
    test_avatar_info()

    # 4. 测试未授权访问
    test_upload_avatar_unauthorized()

    # 5. 测试无效参数
    test_upload_no_file(token)
    test_upload_empty_file(token)

    # 6. 测试无效文件格式
    test_upload_invalid_file_format(token)

    # 7. 测试文件大小限制
    test_upload_large_file(token)

    # 8. 创建基本测试图片
    basic_test_image = "test_avatar_basic.jpg"
    create_test_image(basic_test_image, color="orange")

    # 9. 测试成功上传头像
    test_upload_avatar_success(token, basic_test_image)

    # 10. 测试不同格式
    test_upload_different_formats(token)

    # 11. 测试多次上传
    test_multiple_uploads(token)

    # 12. 测试并发上传
    test_concurrent_uploads(token)

    # 清理基本测试图片
    if os.path.exists(basic_test_image):
        os.remove(basic_test_image)

    print("\n" + "=" * 60)
    print("头像功能测试完成")


if __name__ == "__main__":
    # 确保PIL库可用
    try:
        from PIL import Image
    except ImportError:
        print("❌ 需要安装Pillow库: pip install Pillow")
        exit(1)
    
    run_all_avatar_tests()
